import databricks.koalas as ks
import numpy as np
import pandas as pd
import pyspark.pandas as ps

if __name__ == '__main__':
    kdf = ps.DataFrame({'A': ['foo', 'bar', 'foo', 'bar', 'foo', 'bar', 'foo', 'foo'],
                        'B': ['one', 'one', 'two', 'three', 'two', 'two', 'one', 'three'],
                        'C': np.random.randn(8),
                        'D': np.random.randn(8)})
    kdf.to_csv("../target/koalas/foo.csv")
    print(ps.read_csv("../target/koalas/foo.csv").head(10))

    kdf.to_parquet("../target/koalas/foo.parquet")
    print(ps.read_parquet("../target/koalas/foo.parquet").head(10))

    kdf.to_spark_io("../target/koalas/foo.orc", format="orc")
    print(ps.read_spark_io("../target/koalas/foo.orc", format="orc").head(10))
